package com.test.langchain4j.config;

import com.test.langchain4j.service.ChatAssistant;
import com.test.langchain4j.service.ChatMemoryAssistant;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiTokenCountEstimator;
import dev.langchain4j.service.AiServices;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * Created with IntelliJ IDEA.
 *
 * @description:
 * @author: liuziyang
 * @since: 2025/7/31 14:27
 * @modifiedBy:
 * @version: 1.0
 */
@Configuration
public class LLMConfig {
  @Bean
  public ChatModel chatModel() {
    return OpenAiChatModel.builder()
        .apiKey(System.getenv("qwen-api-key"))
        .modelName("qwen-long")
        .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
        .build();
  }

  @Bean(name = "chat")
  public ChatAssistant chatAssistant(ChatModel chatModel) {
    return AiServices.create(ChatAssistant.class, chatModel);
  }

  @Bean(name = "chatMessageWindowChatMemory")
  public ChatMemoryAssistant chatMessageWindowChatMemory(ChatModel chatModel) {
    return AiServices.builder(ChatMemoryAssistant.class)
        .chatModel(chatModel)
        // 注意每个memoryId对应创建一个ChatMemory
        .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(100))
        .build();
  }

  @Bean(name = "chatTokenWindowChatMemory")
  public ChatMemoryAssistant chatTokenWindowChatMemory(ChatModel chatModel) {
    // TokenCountEstimator默认的token分词器，需要结合Tokenizer计算ChatMessage的token数量
    TokenCountEstimator openAiTokenizer = new OpenAiTokenCountEstimator("gpt-4");

    return AiServices.builder(ChatMemoryAssistant.class)
        .chatModel(chatModel)
        .chatMemoryProvider(memoryId -> TokenWindowChatMemory.withMaxTokens(1000, openAiTokenizer))
        .build();
  }
}
